import asyncio
import time
import numpy as np
import pandas as pd
import sqlite3
from tqsdk import TqApi, TqAuth, TqKq
from MyTT import EMA

class SimpleTradingBot:
    def __init__(self):
        self.api = None
        self.running = True
        self.active_symbols = set()

    async def get_monitored_symbols(self):
        """从数据库获取监控品种"""
        conn = sqlite3.connect('trade_monitor.db')
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM products WHERE is_monitored=1")
        products = cursor.fetchall()
        conn.close()
        # 假设symbol是第3列(索引2)
        return [row[2] for row in products]

    async def run(self):
        """主运行函数"""
        try:
            self.api = TqApi(account=TqKq(), auth=TqAuth("cps168", "alibaba999"))
            await self.api.wait_update()  # 确保添加await
            
            symbols = await self.get_monitored_symbols()
            print(f"开始监控品种: {symbols}")
            
            for symbol in symbols:
                self.api.create_task(self.trade_loop(symbol))
                self.active_symbols.add(symbol)
                
            while self.running:
                await asyncio.sleep(1)
                
        except Exception as e:
            print(f"程序出错: {e}")
        finally:
            while len(self.active_symbols) > 0:
                await asyncio.sleep(0.1)
            
            if hasattr(self, 'api') and self.api is not None:
                try:
                    await self.api.close()  # 这里也需要await
                except Exception as e:
                    print(f"关闭API时出错: {e}")

    async def trade_loop(self, symbol):
        """单个品种的交易循环"""
        try:
            quote = self.api.get_quote(symbol)
            klines = self.api.get_kline_serial(symbol, 60, 500)
            
            while self.running:
                await self.api.wait_update()
                
                # 只有当行情变化时才执行止损检查
                if self.api.is_changing(quote):
                    await self.check_stoploss(symbol, quote, klines)
                    
        except Exception as e:
            print(f"处理品种 {symbol} 时出错: {e}")
        finally:
            self.active_symbols.remove(symbol)

    async def check_stoploss(self, symbol, quote, klines):
        """止损逻辑"""
        print(f"{time.strftime('%X')} {symbol} 价格: {quote.last_price}")
        
        # 计算技术指标
        q = (3*klines.close + klines.low + klines.open + klines.high) / 6
        weights = np.arange(26, 0, -1)
        ma = []
        for i in range(len(q)):
            if i < 25:
                ma.append(np.nan)
                continue
            window = q[i-25:i+1]
            ma.append(np.dot(window, weights) / 351)
        ema_7 = EMA(pd.Series(ma), 7).tolist()
        
        position = self.api.get_position(symbol)
        
        # 平仓条件判断
        if position.pos_short > 0 and ma[-2] > ema_7[-2]:
            await self.api.insert_order(
                symbol=symbol,
                direction="BUY",
                offset="CLOSE",
                volume=position.pos_short
            )
            print(f"{time.strftime('%X')} 平空仓: {symbol} 手数: {position.pos_short}")
            
        if position.pos_long > 0 and ma[-2] < ema_7[-2]:
            await self.api.insert_order(
                symbol=symbol,
                direction="SELL",
                offset="CLOSE",
                volume=position.pos_long
            )
            print(f"{time.strftime('%X')} 平多仓: {symbol} 手数: {position.pos_long}")

if __name__ == "__main__":
    bot = SimpleTradingBot()
    try:
        asyncio.run(bot.run())
    except KeyboardInterrupt:
        bot.running = False
        print("\n程序已停止")
